import numpy as np

# 定义决策树节点类
class Node:
    def __init__(self, predicted_class):
        self.predicted_class = predicted_class  # 每个节点存放的 预测的 类别
        self.feature_index = 0  # 每个节点用来分割的特征的 特征索引
        self.threshold = 0  # 每个节点用来分割特征值的 分割阈值，样本小于则分配到左子树，大于等于则分配到右子树
        self.left = None  # 左子树
        self.right = None  # 右子树

# 定义决策树类
class DecisionTree:
    def __init__(self, max_depth=None, max_features=None, random_state=None):        
        self.max_depth = max_depth  # 决策树的最大深度，用于防止过拟合
        self.max_features = max_features  # 每个节点考虑的最大特征数
        self.random_state = random_state  # 随机数种子
        self.tree = None # 存放构建的决策树

    # 计算基尼系数
    def gini_measure(self, y):
        # counts每类的标签的数量
        unique_classes, counts = np.unique(y, return_counts=True)
        total_samples = len(y) # 标签的总数
        Gini = 1
        for count in counts:
            p = count / total_samples # 计算每类占总数的比例
            Gini -= p ** 2
        return Gini

    # 构建决策树
    def fit(self, X, y):
        self.n_classes = len(set(y)) # 标签的类别总数
        self.n_features = X.shape[1] # 一个样本的特征总数

        if self.max_features is None:
            # Bagging
            # 自举后的数据集使用全部的特征 构建决策树
            self.max_features = self.n_features 

        # 当传入的参数是 浮点型 且 为大于0小于1的小数时
        if 0 < self.max_features <= 1 and isinstance(self.max_features, float):
            # RandomForest
            # 自举后的数据集使用随机抽取的部分特征 构建决策树
            self.max_features = int(self.max_features * self.n_features)

        # 使用递归构建决策树
        self.tree = self.grow_tree(X, y, self.random_state)

    # 根据特征分割数据
    def split_data(self, X, y, random_state):
        # 标签的总数小于1，无法分割，返回None
        m = len(y)
        if m < 1:
            return None
        
        # 计算数据集分割前的基尼系数
        # 如果为0，表示该数据集已经最纯了，无需分割
        impurity_base = self.gini_measure(y)
        if impurity_base == 0:
            return None, None

        # 初始化 用来分割的最优特征 和 最优阈值
        best_index, best_threshold = None, None
        
        # RandomForest
        # 从所有的特征中 随机抽取一定数量的特征
        np.random.seed(random_state)
        # 得到随机抽取后，对应特征的索引
        # cols: fea1 fea2 fea3 fea4 ...
        #  idx:  0    1    2    3   ...
        feature_idxs = np.random.choice(range(self.n_features), self.max_features, replace=False)

        for feature_idx in feature_idxs:
            # 取每个feature那列的数据中的unique数据值
            #     feature_values = [4.35, 4.45, 4.55, 4.65, 4.75, 4.85,]
            feature_values = sorted(set(X[:, feature_idx]))
            # feature_values[1:] = [4.45, 4.55, 4.65, 4.75, 4.85,]
            # 用前后两个数据的均值 作为 可能的分割阈值
            possible_thresholds = [np.mean([i, j]) for i, j in zip(feature_values, feature_values[1:])]
            
            # 依次比较所有的分割阈值，从中找到最优的分割阈值
            # 贪心算法
            for threshold in possible_thresholds:
                # 小于阈值的 特征 对应的标签 为左节点
                left_y = y[X[:, feature_idx] < threshold]
                # 大于阈值的 特征 对应的标签 为右节点
                right_y = y[X[:, feature_idx] >= threshold]
                
                # 左节点 和 右节点 中 标签占 总标签的比例
                left_ratio, right_ratio = len(left_y) / m, len(right_y) / m
                # 计算用 当前该阈值的 分割得到的 基尼指数
                impurity_current = left_ratio * self.gini_measure(left_y) + right_ratio * self.gini_measure(right_y)
               
                # 和上一次 或者 为分割 时的阈值作比较，不断更新，
                # 得到最优的 分割阈值，和该节点 最适合拿来作为分割 特征的索引
                if impurity_current < impurity_base:
                    best_index = feature_idx
                    best_threshold = threshold
                    impurity_base = impurity_current

        return best_index, best_threshold

    # 递归构建决策树
    def grow_tree(self, X, y, random_state, depth=0):
        # self.n_classes 标签的类别总数
        sample_per_class = []
        # 统计每个标签对应的数量
        # 比如 类别 y==0，num of y==0 有4个
        for each_class in range(self.n_classes):
            num_each_class = np.sum(y == each_class)
            sample_per_class.append(num_each_class)
        # 找到出现样本数量最多的那个类别的索引
        predict_class = np.argmax(sample_per_class)
        
        # 将该类别作为根节点
        # 该节点的类别：就是 predict_class
        # self.feature_index = 0  # 特征索引，用于分割数据的特征
        # self.threshold = 0  # 分割阈值，样本小于则分配到左子树，大于等于则分配到右子树
        # self.left = None  # 左子树
        # self.right = None  # 右子树
        node = Node(predicted_class=predict_class)
        '''
        选择出现次数最多的类别作为根节点的原因：
        猜动物游戏示例
        假设我们在一个农场里看到一群动物，动物种类有羊、牛和马（分别对应类别 0、1、2）。
        一共有 6 只动物，其中 1 只是羊，2 只是牛，3 只是马（类似于y = [0, 1, 1, 2, 2, 2]）。
        如果有人问我们这群动物主要是什么，我们会回答是马。
        在决策树的初始节点，这就相当于根据现有的动物数量分布，选择出现次数最多的马这个类别作为预测类别，
        因为从当前看到的情况来说，马是最有可能的主要动物种类。
        '''

        # 使用递归构建决策树
        # 如果最大深度没有设定 或者 当前深度小于设定的最大树深时
        # 不断递归，构建决策树
        if (self.max_depth is None) or (depth < self.max_depth):
            # 得到 最佳的 特征分割索引 和 分割阈值
            idx, threshold = self.split_data(X, y, random_state)
            if idx is not None:
                if random_state is not None:
                    random_state += 1
                # 得到一个 该特征 中特征值 小于 最佳阈值的
                # 一维布尔数组 eg：[True, True, False, False, False]
                # True 的元素 就是 小于 阈值的特征值 的索引
                idx_left = X[:, idx] < threshold
                # 得到小于阈值的特征值
                X_left, y_left = X[idx_left], y[idx_left]
                # False取反得到 大于阈值的特征值
                X_right, y_right = X[~idx_left], y[~idx_left]

                # 将该节点的 特征分割索引 和 分割阈值 
                # 设定未为 最佳的 特征分割索引 和 分割阈值
                node.feature_index = idx
                node.threshold = threshold
                # 使用递归，构建根节点的左右子树
                node.left = self.grow_tree(X_left, y_left, random_state, depth + 1)
                node.right = self.grow_tree(X_right, y_right, random_state, depth + 1)
        return node

    # 对一个样本 预测 类别
    def predict_sample(self, x):
        # 获取决策树的根节点
        node = self.tree  
        # 当前节点左子树不为空，说明还没到叶子节点，继续循环判断
        while node.left:  
            # 如果小于该节点的阈值，就进入左子树节点
            if x[node.feature_index] < node.threshold:  
                node = node.left  
            else:
                # 如果大于等于该节点的阈值，就进入右子树节点
                node = node.right  
        # 为空说明，到叶子节点了，返回该叶子节点的预测类别
        return node.predicted_class  

    # 对数据集进行预测
    def predict(self, X):
        # 从数据集中 取每个样本 进行预测
        # 得到 一个一维数组，包含了所有样本的预测类别
        '''
        input:
                fea1    fea2    fea3    ...
            x1  ...     ...     ...
            x2  ...     ...     ...

        return:
            y = [y1,    y2,     y3      ...]
        '''
        predicted_classes = np.array([self.predict_sample(x) for x in X])
        return predicted_classes